package com.rran.study.chat.gateway.filter;

import com.rran.study.chat.gateway.exception.GatewayErrorEnum;
import com.rran.study.chat.gateway.exception.GatewayException;
import com.rran.study.chat.gateway.limiter.MyRedisRateLimiter;
import com.rran.study.chat.gateway.resolve.ApiKeyResolver;
import org.springframework.boot.context.properties.ConfigurationProperties;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.cloud.gateway.filter.ratelimit.KeyResolver;
import org.springframework.cloud.gateway.filter.ratelimit.RateLimiter;
import org.springframework.cloud.gateway.route.Route;
import org.springframework.cloud.gateway.support.HttpStatusHolder;
import org.springframework.cloud.gateway.support.ServerWebExchangeUtils;
import org.springframework.http.HttpStatus;
import org.springframework.http.MediaType;
import org.springframework.stereotype.Component;

import java.util.Map;

import static org.springframework.cloud.gateway.support.ServerWebExchangeUtils.setResponseStatus;

/**
 * @author yy
 * @Type MyRequestRateLimiterGatewayFilterFactory.java
 * @Desc 自定义限流器
 * @date 2020/8/24
 */
@ConfigurationProperties("spring.cloud.gateway.filter.request-rate-limiter")
@Component
public class MyRequestRateLimiterGatewayFilterFactory extends AbstractGatewayFilterFactory<MyRequestRateLimiterGatewayFilterFactory.Config> {

    public static final String KEY_RESOLVER_KEY = "keyResolver";
    private static final String EMPTY_KEY = "____EMPTY_KEY__";

    private final MyRedisRateLimiter myRedisRateLimiter;
    private final ApiKeyResolver apiKeyResolver;

    /** Switch to deny requests if the Key Resolver returns an empty key, defaults to true. */
    private boolean denyEmptyKey = true;

    /** HttpStatus to return when denyEmptyKey is true, defaults to FORBIDDEN. */
    private String emptyKeyStatusCode = HttpStatus.FORBIDDEN.name();

    public MyRequestRateLimiterGatewayFilterFactory(MyRedisRateLimiter myRedisRateLimiter,
                                                    ApiKeyResolver apiKeyResolver) {
        super(Config.class);
        this.myRedisRateLimiter = myRedisRateLimiter;
        this.apiKeyResolver = apiKeyResolver;
    }

    public ApiKeyResolver getDefaultKeyResolver() {
        return apiKeyResolver;
    }

    public MyRedisRateLimiter getDefaultRateLimiter() {
        return myRedisRateLimiter;
    }

    public boolean isDenyEmptyKey() {
        return denyEmptyKey;
    }

    public void setDenyEmptyKey(boolean denyEmptyKey) {
        this.denyEmptyKey = denyEmptyKey;
    }

    public String getEmptyKeyStatusCode() {
        return emptyKeyStatusCode;
    }

    public void setEmptyKeyStatusCode(String emptyKeyStatusCode) {
        this.emptyKeyStatusCode = emptyKeyStatusCode;
    }

    @SuppressWarnings("unchecked")
    @Override
    public GatewayFilter apply(Config config) {
        KeyResolver resolver = getOrDefault(config.keyResolver, apiKeyResolver);
        RateLimiter<Object> limiter = getOrDefault(config.rateLimiter, myRedisRateLimiter);
        boolean denyEmpty = getOrDefault(config.denyEmptyKey, this.denyEmptyKey);
        HttpStatusHolder emptyKeyStatus = HttpStatusHolder.parse(getOrDefault(config.emptyKeyStatus, this.emptyKeyStatusCode));

        return (exchange, chain) -> {
            Route route = exchange.getAttribute(ServerWebExchangeUtils.GATEWAY_ROUTE_ATTR);

            return resolver.resolve(exchange).defaultIfEmpty(EMPTY_KEY).flatMap(key -> {
                if (EMPTY_KEY.equals(key)) {
                    if (denyEmpty) {
                        setResponseStatus(exchange, emptyKeyStatus);
                        return exchange.getResponse().setComplete();
                    }
                    return chain.filter(exchange);
                }
                return limiter.isAllowed(route.getId(), key).flatMap(response -> {

                    for (Map.Entry<String, String> header : response.getHeaders().entrySet()) {
                        exchange.getResponse().getHeaders().add(header.getKey(), header.getValue());
                    }

                    if (response.isAllowed()) {
                        return chain.filter(exchange);
                    }

                    setResponseStatus(exchange, HttpStatus.TOO_MANY_REQUESTS);
                    exchange.getResponse().getHeaders().set("ContentType", MediaType.APPLICATION_JSON_VALUE);
                    exchange.getResponse().getHeaders().set("DataEncoding","UTF-8");
//                    return exchange.getResponse().setComplete();
                    throw new GatewayException(GatewayErrorEnum.SYSTEM_SERVICE_TOO_MANY_REQUESTS);
                });
            });
        };
    }

    private <T> T getOrDefault(T configValue, T defaultValue) {
        return (configValue != null) ? configValue : defaultValue;
    }

    public static class Config {
        private KeyResolver keyResolver;
        private RateLimiter rateLimiter;
        private HttpStatus statusCode = HttpStatus.TOO_MANY_REQUESTS;
        private Boolean denyEmptyKey;
        private String emptyKeyStatus;

        public KeyResolver getKeyResolver() {
            return keyResolver;
        }

        public Config setKeyResolver(KeyResolver keyResolver) {
            this.keyResolver = keyResolver;
            return this;
        }
        public RateLimiter getRateLimiter() {
            return rateLimiter;
        }

        public Config setRateLimiter(RateLimiter rateLimiter) {
            this.rateLimiter = rateLimiter;
            return this;
        }

        public HttpStatus getStatusCode() {
            return statusCode;
        }

        public Config setStatusCode(HttpStatus statusCode) {
            this.statusCode = statusCode;
            return this;
        }

        public Boolean getDenyEmptyKey() {
            return denyEmptyKey;
        }

        public Config setDenyEmptyKey(Boolean denyEmptyKey) {
            this.denyEmptyKey = denyEmptyKey;
            return this;
        }

        public String getEmptyKeyStatus() {
            return emptyKeyStatus;
        }

        public Config setEmptyKeyStatus(String emptyKeyStatus) {
            this.emptyKeyStatus = emptyKeyStatus;
            return this;
        }
    }

}